package handlers

import (
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/gorilla/mux"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源
	},
}

// WSMessage WebSocket消息
type WSMessage struct {
	Type string      `json:"type"`
	Data interface{} `json:"data"`
}

// WSHub WebSocket集线器
type WSHub struct {
	clients    map[*websocket.Conn]bool
	broadcast  chan WSMessage
	register   chan *websocket.Conn
	unregister chan *websocket.Conn
	mutex      sync.RWMutex
}

// NewWSHub 创建WebSocket集线器
func NewWSHub() *WSHub {
	return &WSHub{
		clients:    make(map[*websocket.Conn]bool),
		broadcast:  make(chan WSMessage),
		register:   make(chan *websocket.Conn),
		unregister: make(chan *websocket.Conn),
	}
}

// Run 运行WebSocket集线器
func (h *WSHub) Run() {
	for {
		select {
		case conn := <-h.register:
			h.mutex.Lock()
			h.clients[conn] = true
			h.mutex.Unlock()
			log.Printf("WebSocket client connected: %v", conn.RemoteAddr())

		case conn := <-h.unregister:
			h.mutex.Lock()
			if _, ok := h.clients[conn]; ok {
				delete(h.clients, conn)
				conn.Close()
			}
			h.mutex.Unlock()
			log.Printf("WebSocket client disconnected: %v", conn.RemoteAddr())

		case message := <-h.broadcast:
			h.mutex.RLock()
			for conn := range h.clients {
				select {
				case <-time.After(10 * time.Second):
					// 写入超时，关闭连接
					delete(h.clients, conn)
					conn.Close()
				default:
					if err := conn.WriteJSON(message); err != nil {
						delete(h.clients, conn)
						conn.Close()
					}
				}
			}
			h.mutex.RUnlock()
		}
	}
}

// Broadcast 广播消息
func (h *WSHub) Broadcast(msgType string, data interface{}) {
	select {
	case h.broadcast <- WSMessage{Type: msgType, Data: data}:
	default:
		// 广播队列满，跳过消息
	}
}

// GetClientCount 获取连接客户端数量
func (h *WSHub) GetClientCount() int {
	h.mutex.RLock()
	defer h.mutex.RUnlock()
	return len(h.clients)
}

// WSHandler WebSocket处理器
type WSHandler struct {
	hub *WSHub
}

// NewWSHandler 创建WebSocket处理器
func NewWSHandler() *WSHandler {
	hub := NewWSHub()
	go hub.Run()
	return &WSHandler{hub: hub}
}

// HandleWebSocket 处理WebSocket连接
func (h *WSHandler) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return
	}
	defer conn.Close()

	// 注册客户端
	h.hub.register <- conn

	// 发送欢迎消息
	welcome := WSMessage{
		Type: "connected",
		Data: map[string]interface{}{
			"message": "WebSocket connected successfully",
			"time":    time.Now(),
		},
	}
	conn.WriteJSON(welcome)

	// 读取客户端消息（保持连接活跃）
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("WebSocket error: %v", err)
			}
			break
		}
	}

	// 注销客户端
	h.hub.unregister <- conn
}

// BroadcastTaskEvent 广播任务事件
func (h *WSHandler) BroadcastTaskEvent(eventType string, task interface{}) {
	h.hub.Broadcast("task_event", map[string]interface{}{
		"event": eventType,
		"task":  task,
		"time":  time.Now(),
	})
}

// BroadcastLogEvent 广播日志事件
func (h *WSHandler) BroadcastLogEvent(log interface{}) {
	h.hub.Broadcast("log_event", map[string]interface{}{
		"log":  log,
		"time": time.Now(),
	})
}

// BroadcastSystemEvent 广播系统事件
func (h *WSHandler) BroadcastSystemEvent(event string, data interface{}) {
	h.hub.Broadcast("system_event", map[string]interface{}{
		"event": event,
		"data":  data,
		"time":  time.Now(),
	})
}

// GetHub 获取WebSocket集线器
func (h *WSHandler) GetHub() *WSHub {
	return h.hub
}

// SetupWSRoutes 设置WebSocket路由
func SetupWSRoutes(router *mux.Router, wsHandler *WSHandler) {
	router.HandleFunc("/ws", wsHandler.HandleWebSocket)
} 